from clusters_generate import *
import numpy as np
import matplotlib.pyplot as plt
from plyfile import PlyData, PlyElement
import math
from sklearn.cluster import DBSCAN
from sklearn.neighbors import NearestNeighbors

def auto_tune_parameters(points):
    # 计算平均最近邻距离
    nbrs = NearestNeighbors(n_neighbors=5).fit(points)
    distances, _ = nbrs.kneighbors(points)
    mean_distance = np.mean(distances[:,-1])
    
    # 动态设置参数
    auto_eps = mean_distance * 3
    auto_min_samples = int(len(points)*0.05)
    
    return DBSCAN(eps=auto_eps, min_samples=auto_min_samples)

def calculate_clusters(vertices):
    points = np.array([[v['x'], v['y'], v['z']] for v in vertices])
    
    # 使用DBSCAN进行聚类
    db = auto_tune_parameters(points).fit(points)
    labels = db.labels_
    
    clusters = []
    
    # 遍历所有簇（忽略噪声点-1）
    for cluster_id in set(labels) - {-1}:
        # 获取当前簇的所有点
        cluster_points = points[labels == cluster_id]
        
        # 计算中心点（质心）
        center = cluster_points.mean(axis=0)
        
        # 计算半径（最大距离）
        distances = np.linalg.norm(cluster_points - center, axis=1)
        radius = distances.max()
        
        clusters.append({
            'center': center.tolist(),
            'radius': float(radius)
        })
    
    # 在返回clusters前添加过滤逻辑
    filtered_clusters = []
    for i, cluster in enumerate(clusters):
        is_nested = False
        for j, other in enumerate(clusters):
            if i != j:
                # 计算簇间距离
                distance = np.linalg.norm(np.array(cluster['center']) - np.array(other['center']))
                # 如果当前簇完全被其他簇包含
                if distance + cluster['radius'] < other['radius']:
                    is_nested = True
                    break
        if not is_nested:
            filtered_clusters.append(cluster)
            
    return filtered_clusters

def detect_clusters(filename):
        # 读取PLY文件
    ply_data = PlyData.read(filename)
    vertices = ply_data['vertex']
    
    # 提取坐标数据
    x = vertices['x']
    y = vertices['y']
    z = vertices['z']

    # 计算深度范围
    min_depth = np.min(z)
    max_depth = np.max(z)
    depth_range = max_depth - min_depth

    colors = []
    for depth in z:
        # 将深度映射到颜色索引 (0-255)
        color_idx = int(((depth - min_depth) / depth_range) * (len(color_map)-1))
        colors.append(color_map[color_idx])
    
    # 转换为0-1范围的RGB值
    colors = np.array(colors) / 255.0

    # 创建3D可视化
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制原始点云
    ax.scatter(x, y, z, c=colors, s=1, marker='o', label='原始点云')

    clusters = calculate_clusters(vertices)
    print(clusters)

    # 绘制聚类结果
    for cluster in clusters:
        center = cluster['center']
        radius = cluster['radius']
        # 绘制 sphere
        u, v = np.mgrid[0:2*np.pi:20j, 0:np.pi:10j]
        x = center[0] + radius * np.cos(u) * np.sin(v)
        y = center[1] + radius * np.sin(u) * np.sin(v)
        z = center[2] + radius * np.cos(v)
        ax.plot_surface(x, y, z, color='b', alpha=0.3)

    # 设置坐标轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    
    plt.show(block=True)  # 添加block=True参数保持窗口显示

if __name__ == "__main__":
    detect_clusters('clusters.ply')
